use std::{marker::PhantomData, ops::Deref};

use wgpu::util::DeviceExt;

use crate::{material::Attribute, AsBytes, Context, Pixmap};

pub struct DynBuffer<D: AsBytes> {
    buffer: wgpu::Buffer,
    _p: PhantomData<D>,
}

impl<D: AsBytes> DynBuffer<D> {
    pub fn new(ctx: &Context, amount: usize, usage: wgpu::BufferUsages) -> Self {
        let buffer = ctx.device.create_buffer(&wgpu::BufferDescriptor {
            label: None,
            size: (amount * std::mem::size_of::<D>()) as _,
            mapped_at_creation: false,
            usage: usage | wgpu::BufferUsages::COPY_SRC | wgpu::BufferUsages::COPY_DST,
        });

        Self {
            buffer,
            _p: PhantomData,
        }
    }

    pub fn write(&mut self, ctx: &super::Context, data: &[D]) {
        let bytes = data.as_bytes();
        if bytes.len() as u64 > self.buffer.size() {
            self.resize(ctx, data.len())
        }

        ctx.queue.write_buffer(&self.buffer, 0, bytes)
    }

    pub fn resize(&mut self, ctx: &super::Context, amount: usize) {
        let size = (amount * std::mem::size_of::<D>()) as _;
        if size <= self.buffer.size() {
            return;
        }

        let new_buffer = ctx.device.create_buffer(&wgpu::BufferDescriptor {
            label: None,
            size,
            mapped_at_creation: false,
            usage: self.buffer.usage(),
        });

        if self.buffer.size() > 0 {
            let mut cmd = ctx.device.create_command_encoder(&Default::default());
            cmd.copy_buffer_to_buffer(&self.buffer, 0, &new_buffer, 0, self.buffer.size());
            ctx.queue.submit(std::iter::once(cmd.finish()));
        }

        self.buffer = new_buffer;
    }
}

impl<D: AsBytes> Deref for DynBuffer<D> {
    type Target = wgpu::Buffer;

    fn deref(&self) -> &Self::Target {
        &self.buffer
    }
}

pub struct DeviceTexture2D {
    texture: wgpu::Texture,
    view: wgpu::TextureView,
}

impl DeviceTexture2D {
    pub fn new(ctx: &Context, width: u32, height: u32, format: wgpu::TextureFormat) -> Self {
        Self::with_extent(
            ctx,
            wgpu::Extent3d {
                width,
                height,
                depth_or_array_layers: 1,
            },
            format,
        )
    }

    pub fn with_extent(ctx: &Context, extent: wgpu::Extent3d, format: wgpu::TextureFormat) -> Self {
        let texture = ctx.device.create_texture(&wgpu::TextureDescriptor {
            label: None,
            size: extent,
            mip_level_count: 1,
            sample_count: 1,
            dimension: wgpu::TextureDimension::D2,
            format,
            usage: wgpu::TextureUsages::TEXTURE_BINDING | wgpu::TextureUsages::COPY_DST,
            view_formats: &[],
        });

        let view = texture.create_view(&Default::default());

        Self { texture, view }
    }

    pub fn with_pixmap(ctx: &Context, pixmap: &Pixmap, format: wgpu::TextureFormat) -> Self {
        let texture = ctx.device.create_texture_with_data(
            &ctx.queue,
            &wgpu::TextureDescriptor {
                label: None,
                size: pixmap.extent(),
                mip_level_count: 1,
                sample_count: 1,
                dimension: wgpu::TextureDimension::D2,
                format,
                usage: wgpu::TextureUsages::TEXTURE_BINDING,
                view_formats: &[],
            },
            wgpu::util::TextureDataOrder::MipMajor,
            pixmap.pixels.as_slice().as_bytes(),
        );

        let view = texture.create_view(&Default::default());

        Self { texture, view }
    }

    pub fn write_pixels(
        &self,
        ctx: &Context,
        x: u32,
        y: u32,
        width: u32,
        height: u32,
        data: &[u8],
        data_layout: wgpu::ImageDataLayout,
    ) {
        ctx.queue.write_texture(
            wgpu::ImageCopyTexture {
                texture: &self.texture,
                mip_level: 0,
                origin: wgpu::Origin3d { x, y, z: 0 },
                aspect: wgpu::TextureAspect::All,
            },
            data,
            data_layout,
            wgpu::Extent3d {
                width,
                height,
                depth_or_array_layers: 1,
            },
        )
    }

    pub fn resize(&mut self, ctx: &Context, width: u32, height: u32) {
        let format = self.texture.format();
        let mut size = self.texture.size();
        size.width = width;
        size.height = height;
        self.texture = ctx.device.create_texture(&wgpu::TextureDescriptor {
            label: None,
            size,
            mip_level_count: 1,
            sample_count: 1,
            dimension: wgpu::TextureDimension::D2,
            format,
            usage: wgpu::TextureUsages::TEXTURE_BINDING | wgpu::TextureUsages::COPY_DST,
            view_formats: &[],
        });
        self.view = self.texture.create_view(&Default::default());
    }

    pub fn raw_texture(&self) -> &wgpu::Texture {
        &self.texture
    }
}

impl Attribute for DeviceTexture2D {
    fn ty(&self) -> wgpu::BindingType {
        wgpu::BindingType::Texture {
            sample_type: match self.format() {
                wgpu::TextureFormat::R8Uint => wgpu::TextureSampleType::Uint,
                _ => wgpu::TextureSampleType::Float { filterable: true },
            },
            view_dimension: wgpu::TextureViewDimension::D2,
            multisampled: false,
        }
    }

    fn res(&self) -> wgpu::BindingResource {
        wgpu::BindingResource::TextureView(&self.view)
    }
    
    fn prepare(&mut self, _: &Context) {
        
    }
}

impl Deref for DeviceTexture2D {
    type Target = wgpu::Texture;

    fn deref(&self) -> &Self::Target {
        &self.texture
    }
}
